/*
 * Copyright (c) 2010-2025 Haifeng Li. All rights reserved.
 *
 * Smile is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Smile is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Smile. If not, see <https://www.gnu.org/licenses/>.
 */
package smile.plot.vega;

import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;

/**
 * The regression transform fits two-dimensional regression models to smooth
 * and predict data. This transform can fit multiple models for input data
 * (one per group) and generates new data objects that represent points for
 * summary trend lines. Alternatively, this transform can be used to generate
 * a set of objects containing regression model parameters, one per group.
 * <p>
 * This transform supports parametric models for the following functional forms:
 * <p>
 * linear (linear): y = a + b * x
 * logarithmic (log): y = a + b * log(x)
 * exponential (exp): y = a * e^(b * x)
 * power (pow): y = a * x^b
 * quadratic (quad): y = a + b * x + c * x^2
 * polynomial (poly): y = a + b * x + … + k * x^(order)
 * <p>
 * All models are fit using ordinary least squares. For non-parametric locally
 * weighted regression, see the loess transform.
 *
 * @author Haifeng Li
 */
public class RegressionTransform {
    /** VegaLite's Regression definition object. */
    final ObjectNode spec;

    /**
     * Hides the constructor so that users cannot create the instances directly.
     */
    RegressionTransform(ObjectNode spec) {
        this.spec = spec;
    }

    @Override
    public String toString() {
        return spec.toString();
    }

    /**
     * Returns the specification in pretty print.
     * @return the specification in pretty print.
     */
    public String toPrettyString() {
        return spec.toPrettyString();
    }

    /**
     * Sets the functional form of the regression model.
     * @param method "linear", "log", "exp", "pow", "quad", or "poly".
     * @return this object.
     */
    public RegressionTransform method(String method) {
        spec.put("method", method);
        return this;
    }

    /**
     * Sets if the transform should return the regression model parameters
     * (one object per group), rather than trend line points. The resulting
     * objects include a coef array of fitted coefficient values (starting
     * with the intercept term and then including terms of increasing order)
     * and an rSquared value (indicating the total variance explained by the
     * model).
     *
     * @param flag a flag indicating if the transform should return the
     *            regression model parameters.
     * @return this object.
     */
    public RegressionTransform params(boolean flag) {
        spec.put("params", flag);
        return this;
    }

    /**
     * Sets the polynomial order (number of coefficients) for the "poly" method.
     *
     * @param order The polynomial order (number of coefficients).
     * @return this object.
     */
    public RegressionTransform order(int order) {
        spec.put("order", order);
        return this;
    }

    /**
     * Sets a [min, max] domain over the independent (x) field for the
     * starting and ending points of the generated trend line.
     * @param min the minimum value of the independent field domain.
     * @param max the maximum value of the independent field domain.
     * @return this object.
     */
    public RegressionTransform extent(double min, double max) {
        spec.putArray("extent").add(min).add(max);
        return this;
    }

    /**
     * Sets the data fields to group by. If not specified, a single group
     * containing all data objects will be used.
     *
     * @param fields The data fields to group by. If not specified,
     *              a single group containing all data objects will be used.
     * @return this object.
     */
    public RegressionTransform groupby(String... fields) {
        ArrayNode node = spec.putArray("groupby");
        for (var field : fields) {
            node.add(field);
        }
        return this;
    }

    /**
     * Sets the output field names for the smoothed points generated
     * by the loess transform.
     *
     * @param fields The output field names for the smoothed points
     *              generated by the loess transform.
     * @return this object.
     */
    public RegressionTransform as(String... fields) {
        ArrayNode node = spec.putArray("as");
        for (var field : fields) {
            node.add(field);
        }
        return this;
    }
}
